///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/er_local.cpp
///
/// \brief   implementation file for Edmiston-Ruedenberg orbital localization
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<string>
#include<vector>
#include<complex>

#include<formic/ham/er_local.h>
#include<formic/lbfgs.h>
#include<formic/ham/ham.h>
#include<formic/ham/integral_transform.h>
#include<formic/input/base.h>
#include<formic/exception.h>
#include<formic/output/output.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/lapack/interface.h>
#include<formic/all_data.h>
#include<formic/archive.h>
#include<formic/random/random.h>
#include<formic/numeric/matrix.h>
#include<formic/vector/matrix.h>

namespace formic {

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  class used to perform Edmiston-Ruedenberg localizations on selected groups of orbitals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
class ER_Localize_Worker {

  private:

    /// \brief  a convenient typedef for complex numbers
    typedef std::complex<double> _cmplx;

    /// \brief  the number of orbitals in the orbital set we are localizing
    int _n;

    /// \brief  the maximum number of iterations
    int _maxiter;

    /// \brief  the convergence threshold
    double _thresh;

    /// \brief  the number zero in complex number form
    _cmplx _zcmplx;

    /// \brief  un-modified two-electron integrals for the orbital set we are localizing
    std::vector<_cmplx> _umtei;

    /// \brief  the two-electron integrals for the orbital set we are localizing
    std::vector<_cmplx> _tei;

    /// \brief  the cumulative orbital transformation
    std::vector<double> _cot;

    /// \brief  the orbital transformation ("rotation") matrix computed for a given epsilon
    std::vector<_cmplx> _rot;

    /// \brief  the antisymmetric transformation matrix that when exponentiated with a small
    ///         weight is guaranteed to give a unitary matrix that increases the localization
    std::vector<_cmplx> _astm;

    /// \brief  eigenvalues of the antisymmetric transformation matrix
    std::vector<_cmplx> _astm_evals;

    /// \brief  unitary matrix of right eigenvectors that diagonalize the antisymmetric transformation matrix
    std::vector<_cmplx> _astm_evecs;

    /// \brief  exponential of antisymmetric transformation matrix eigenvalues times rotation size
    std::vector<_cmplx> _exp_evals;

    /// \brief  the two-electron integrals transformed by the eigenvectors of the antisymmetric transformation matrix
    std::vector<_cmplx> _aetei;

    /// \brief  complex work space
    std::vector<_cmplx> _cws;

    /// \brief  real work space
    std::vector<double> _rws;

  private:

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  whether the code uses direct matrix exponentiation to compute the rotation matrix
    ///         from the antisymmetric transformation matrix
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    static bool direct_matrix_exponentiation() { return true; }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  constructs the antisymmetric transformation matrix
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void build_astm() {

      for (int i = 0; i < _n; i++)
      for (int j = 0; j < _n; j++)
        _astm[j+i*_n] = _tei[ _n*_n*_n*i + _n*_n*i + _n*i + j ] - _tei[ _n*_n*_n*i + _n*_n*j + _n*j + j ];

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  diagonalizes the antisymmetric transformation matrix
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void diagonalize_astm() {

      // convert the matrix to complex numbers
      for (int i = 0; i < _n*_n; i++)
        _cws[i] = formic::unity(_cmplx()) * _astm[i];

      // diagonalize, placing the eigenvalues in _astm_evals and the eigenvectors in the columns of _astm_evecs
      int info;
      int lwork = 4*_n;
      formic::zgeev('N', 'V', _n, &_cws.at(0), _n, &_astm_evals.at(0), 0, _n, &_astm_evecs.at(0), _n, &_cws.at(_n*_n), lwork, &_rws.at(0), info);
      if ( info != 0 )
        throw formic::Exception("formic::zgeev failed for  antisymmetric transformation matrix with info = %i in formic::ER_Localize_Worker::diagonalize_astm") % info;

      // normalize the eigenvectors
      for (int i = 0; i < _n; i++)
        formic::xscal(_n, formic::unity(_cmplx()) / std::abs( std::sqrt( formic::xdotc(_n, &_astm_evecs.at(_n*i), 1, &_astm_evecs.at(_n*i), 1) ) ), &_astm_evecs.at(_n*i), 1);

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  prepare transformation of two electron integrals by antisymmetric transformation
    ///         matrix eigenvectors
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void transform_tei_by_astm_evecs() {

      // get transpose of evec matrix
      _cmplx * tevec = &_cws.at(_n*_n*_n*_n);
      for (int i = 0; i < _n; i++)
      for (int j = 0; j < _n; j++)
        tevec[i+j*_n] = _astm_evecs[j+i*_n];

      // first transform of tei via the evecs
      formic::transform_dense_tei(_n, tevec, tevec, &_tei.at(0), &_aetei.at(0), &_cws.at(0));

      // second transform of tei via the evecs
      for (int p = 0; p < _n; p++)
      for (int q = 0; q < _n; q++)
      for (int r = 0; r < _n; r++)
      for (int s = 0; s < _n; s++) {
        _cmplx d = _zcmplx;
        for (int n = 0; n < _n; n++)
          d += _astm_evecs[n+p*_n] * _astm_evecs[n+q*_n] * _astm_evecs[n+r*_n] * _astm_evecs[n+s*_n];
        _aetei[ _n*_n*_n*p + _n*_n*q + _n*r + s ] *= std::conj(d);
      }

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  exponentiate the antisymmetric transformation matrix's eigenvalues for the supplied
    ///         value of the epsilon parameter
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void exponentiate_evals(const double eps) {

      for (int i = 0; i < _n; i++)
        _exp_evals.at(i) = std::exp( eps * _astm_evals.at(i) );

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  computes the new orbital rotation matrix for the supplied value of the epsilon
    ///         parameter
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void compute_rotation(const double eps) {

      if ( this->direct_matrix_exponentiation() ) {

        // build a matrix around the rotation matrix data array
        formic::Matrix<_cmplx,int> rmat(_n, _n, &_rot.at(0));

        // compute exponential of scaled antisymmetric transformation matrix
        rmat <<= formic::matrix_exponent( -eps * formic::unity(_cmplx()) * formic::ConstMatrix<_cmplx,int>(_n, _n, &_astm.at(0)), 1.0e-12 );

      } else {

        // exponentiate the antisymmetric transformation matrix's eigenvalues
        this->exponentiate_evals(eps);

        // prepare scaled eigenvectors
        for (int j = 0; j < _n; j++)
        for (int i = 0; i < _n; i++)
          _cws[j+i*_n] = _exp_evals[j] * _astm_evecs[i+j*_n];

        // prepare complex conjugates of eigenvectors
        for (int j = 0; j < _n; j++)
        for (int i = 0; i < _n; i++)
          _cws[_n*_n + i+j*_n] = formic::conj(_astm_evecs[i+j*_n]);

        // compute rotation matrix and store it in _rot
        formic::xgemm('N', 'N', _n, _n, _n, formic::unity(_cmplx()), &_cws.at(_n*_n), _n, &_cws.at(0), _n, formic::zero(_cmplx()), &_rot.at(0), _n);

      }

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  get the new function value for a given value of the epsilon parameter
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    double get_value_for_epsilon(const double eps) {

      // initialize return value
      _cmplx retval = _zcmplx;

      if ( this->direct_matrix_exponentiation() ) {

        this->compute_rotation(eps);
        formic::transform_dense_tei(_n, &_rot.at(0), &_rot.at(0), &_tei.at(0), &_aetei.at(0), &_cws.at(0));
        for (int p = 0; p < _n; p++)
          retval += _aetei[ _n*_n*_n*p + _n*_n*p + _n*p + p ];

      } else {

        // exponentiate the antisymmetric transformation matrix's eigenvalues
        this->exponentiate_evals(eps);

        // compute return value
        for (int p = 0; p < _n; p++)
        for (int q = 0; q < _n; q++)
        for (int r = 0; r < _n; r++)
        for (int s = 0; s < _n; s++)
          retval += _aetei[ _n*_n*_n*p + _n*_n*q + _n*r + s ] * _exp_evals[p] * _exp_evals[q] * _exp_evals[r] * _exp_evals[s];

        // ensure complex part is small
        if ( std::abs( formic::imag(retval) ) > 1.0e-06 )
          throw formic::Exception("unexpectedly large imaginary component in formic::ER_Localize_Worker::get_value_for_epsilon");

      }

      // return the value
      return formic::real(retval);

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   transforms two electron integrals to the current orbtial basis
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void transform_ints_by_cot() {

      // get orbitals in complex format
      for (int i = 0; i < _n*_n; i++)
        _cws[ _n*_n*_n*_n + i ] = formic::unity(_cmplx()) * _cot[i];

      // get orbital complex conjugates in complex format
      for (int i = 0; i < _n*_n; i++)
        _cws[ _n*_n*_n*_n + _n*_n + i ] = formic::unity(_cmplx()) * formic::conj(_cot[i]);

      // transform two electron integrals
      formic::transform_dense_tei(_n, &_cws.at(_n*_n*_n*_n), &_cws.at(_n*_n*_n*_n + _n*_n), &_umtei.at(0), &_tei.at(0), &_cws.at(0));

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   orthonormalizes the cumulative orbital transformation matrix via S^(-1/2)
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void orthonormalize_cot() {

      // compute the s_neg_half matrix and store it in the _rws array
      formic::xgemm('N', 'C', _n, _n, _n, 1.0, &_cot.at(0), _n, &_cot.at(0), _n, 0.0, &_rws.at(0), _n);
      formic::matrix_neg_half(_n, &_rws.at(0));

      // orthonormalize the cumulative orbital transformation
      formic::xcopy(_n*_n, &_cot.at(0), 1, &_rws.at(_n*_n), 1);
      formic::xgemm('N', 'N', _n, _n, _n, 1.0, &_rws.at(0), _n, &_rws.at(_n*_n), _n, 0.0, &_cot.at(0), _n);

    }

  public:

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   constructor for the Edmiston-Ruedenberg localization worker
    ///
    /// \param[in]      userinp        user input options
    /// \param[in]      ham            the hamiltonian, which holds the integrals
    /// \param[in]      orb_inds       indices of the orbitals in the set we are localizing
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    ER_Localize_Worker(const formic::InputBase & userinp,
                       const formic::Hamiltonian<double> & ham,
                       const std::vector<int> & orb_inds)
      : _n(orb_inds.size()),
        _maxiter(userinp.get<int>("er_max_iter")),
        _thresh(userinp.get<double>("er_thresh")),
        _zcmplx( formic::zero(_cmplx()) )
    {

      // check that the orbital indices are all different
      if ( ! formic::all_different_by_loops(orb_inds.begin(), orb_inds.end()) )
        throw formic::Exception("repeated orbital index in constructor for ER_Localize_Worker");

      // allocate space for internal arrays
      _umtei.assign(_n*_n*_n*_n, _zcmplx);
      _tei.assign(_n*_n*_n*_n, _zcmplx);
      _cot.assign(_n*_n, 0.0);
      _rot.assign(_n*_n, _zcmplx);
      _astm.assign(_n*_n, _zcmplx);
      _astm_evals.assign(_n, _zcmplx);
      _astm_evecs.assign(_n*_n, _zcmplx);
      _exp_evals.assign(_n, _zcmplx);
      _aetei.assign(_n*_n*_n*_n, _zcmplx);
      _cws.assign(std::max(_n*_n*_n*_n+2*_n*_n, 2*_n*_n+4*_n), _zcmplx);
      _rws.assign(std::max(2*_n*_n, 2*_n), 0.0);

      // initialize two electron integral array for our set of orbitals
      for (int i = 0; i < _n; i++) {
        const int ii = orb_inds.at(i);
        for (int j = 0; j < _n; j++) {
          const int jj = orb_inds.at(j);
          for (int k = 0; k < _n; k++) {
            const int kk = orb_inds.at(k);
            for (int l = 0; l < _n; l++) {
              const int ll = orb_inds.at(l);
              _umtei.at( _n*_n*_n*i + _n*_n*j + _n*k + l ) = formic::unity(_cmplx()) * ham.get_tei(ii,jj,kk,ll);
            }
          }
        }
      }

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   computes the local orbitals
    ///
    /// \param[in,out]  new_orbs       on entry, the initial guess for the local orbitals
    ///                                on exit, the local orbitals
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void localize(std::vector<double> & new_orbs) {

      // check that the orbital coefficient matrix is the right size
      if ( new_orbs.size() != _n*_n )
        throw formic::Exception("Size of new_orbs array (%i) is incorrect in ER_Localize_Worker::localize.  Expected size of %i.") % new_orbs.size() % (_n*_n);

      // set the cumulative transformation to the initial guess
      _cot = new_orbs;

      // orthonormalize the cumulative transformation
      this->orthonormalize_cot();

      // transform integrals based on the initial guess
      this->transform_ints_by_cot();

      // iteratively localize the orbitals
      double best_val = -1.0;
      for (int iter = 0; iter < _maxiter; iter++) {

        // print a greeting
        formic::of << boost::format("starting ER localization iteration %2i") % iter << std::endl << std::endl;

        // build the antisymmetric transformation matrix
        this->build_astm();

        // if we are not directly exponentiating the matrix, then we need to diagonalize it and transform the integrals by its eigenvectors
        if ( ! this->direct_matrix_exponentiation() ) {
          this->diagonalize_astm();
          this->transform_tei_by_astm_evecs();
        }

        // get the initial target function value
        const double curr_val = this->get_value_for_epsilon(0.0);

        // check for convergence
        if ( std::abs( best_val - curr_val ) < _thresh ) {
          formic::of << boost::format("ER localization converged") << std::endl << std::endl;
          break;
        }

        // remember the best value
        best_val = std::max(best_val, curr_val);

        // search for the best magnitude epsilon for the transformation
        double eps_init = 0.0;
        double eps_best = 1.0;
        double search_direction = 1.0;
        double step_size = 1.0;
        formic::lbfgs_loose_line_search(*this, 1, &eps_init, &search_direction, &eps_best, step_size, false);

        // compute resulting orbital transformation and store it in _rot
        this->compute_rotation(eps_best);

        // print rotation matrix
        formic::of << "new rotation matrix:" << std::endl;
        for (int j = 0; j < _n; j++)
        for (int i = 0; i < _n; i++)
          _rws[i+j*_n] = formic::real(_rot[i+j*_n]);
        this->print(_rws);
        formic::of << std::endl;

        // update the cumulative transformation matrix
        std::fill(_rws.begin(), _rws.end(), 0.0);
        for (int k = 0; k < _n; k++)
        for (int j = 0; j < _n; j++)
        for (int i = 0; i < _n; i++)
          _rws[i+j*_n] += formic::real(_rot[i+k*_n]) * _cot[k+j*_n];
        formic::xcopy(_n*_n, &_rws.at(0), 1, &_cot.at(0), 1);

        // orthonormalize the cumulative transformation
        this->orthonormalize_cot();

        // transform the integrals
        this->transform_ints_by_cot();

        // print the cumulative transformation matrix
        formic::of << "cumulative orbital transformation:" << std::endl;
        this->print(_cot);
        formic::of << std::endl;

      }

      // return the cumulative orbital transformation
      new_orbs = _cot;

    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   computes the ER target function value for the line search
    ///
    /// \param[in]      n      number of parameters (one in this case)
    /// \param[in]      eps    pointer to the value of epsilon, the "rotation magnitude"
    /// \param[out]     val    on exit, the ER target function value for the given eps
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void get_value(const int n, const double * const eps, double & val) {
      val = -this->get_value_for_epsilon(*eps);
    }

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief   print out the supplied coefficient matrix
    ///
    /// \param[in]      orbs   vector of orbital coefficients (stored in rows, but printed in columns)
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    void print(const std::vector<double> & orbs) {
      for (int p = 0; p < _n; p++) {
        for (int i = 0; i < _n; i++)
          formic::of << boost::format(" %16.8f") % orbs[i+p*_n];
        formic::of << std::endl;
      }
    }

};

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   performs an Edmiston-Ruedenberg localization on the chosen set of orbitals
///
/// \param[in]      userinp        user input options
/// \param[in,out]  ham            The hamiltonian, which holds the integrals.
///                                On exit, ham will be transformed to the localized basis.
/// \param[in]      orb_inds       indices of orbitals that will be replaced
/// \param[in]      coeff_guess    initial guess for the orbital coefficients
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::er_localize(const formic::InputBase & userinp,
                         formic::Hamiltonian<double> & ham,
                         const std::vector<int> & orb_inds,
                         const std::vector<double> & coeff_guess) {

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // get number of orbitals to localize
  const int nloc = orb_inds.size();

  // print some info
  if ( formic::mpi::rank() == 0 ) {
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << boost::format("   ======   Edmiston Ruedenberg Localization   ======") << std::endl;
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << std::endl;
    formic::of << boost::format("localizing orbitals ");
    for (int i = 0; i < nloc; i++)
      formic::of << boost::format(" %2i") % orb_inds.at(i);
    formic::of << std::endl;
    formic::of << std::endl;
  }

  // allocate the total orbital transformation matrix (includes non-localized orbitals)
  std::vector<double> all_orbs(nao*nao, 0.0);

  // compute orbital localization on root process
  if ( formic::mpi::rank() == 0 ) {

    // create the worker for running the ER algorithm
    formic::ER_Localize_Worker worker(userinp, ham, orb_inds);

    // prepare initial guess for new orbitals
    std::vector<double> new_orbs(nloc*nloc, 0.0);
    if ( coeff_guess.size() != nloc*nloc )
      throw formic::Exception("orbital coefficient guess is the wrong size in formic::er_localize");
    for (int i = 0; i < nloc; i++)
    for (int j = 0; j < nloc; j++)
      new_orbs.at(i+j*nloc) = coeff_guess.at(i+j*nloc);

    // compute localized orbitals
    worker.localize(new_orbs);

    // get the set of unchanged orbitals
    const std::vector<int> unc_orb_inds = formic::other_indices(nao, orb_inds);
    //std::set<int> rep_orb_set(orb_inds.begin(), orb_inds.end());
    //std::set<int> unc_orb_set;
    //for (int j = 0; j < nao; j++)
    //  if ( rep_orb_set.count(j) == 0 )
    //    unc_orb_set.insert(j);
    //std::vector<int> unc_orb_inds(unc_orb_set.begin(), unc_orb_set.end());

    // compute total orbital transformation matrix (including non-localized orbitals)
    for (int i = 0; i < unc_orb_inds.size(); i++)
      all_orbs.at( unc_orb_inds.at(i) + unc_orb_inds.at(i) * nao ) = 1.0;
    for (int j = 0; j < orb_inds.size(); j++)
    for (int i = 0; i < orb_inds.size(); i++) {
      const int ii = orb_inds.at(i);
      const int jj = orb_inds.at(j);
      all_orbs.at( ii + jj * nao ) = new_orbs.at(i+j*nloc);
    }

  }

  // perform the orbital transformation
  ham.transform_orbitals(&all_orbs.at(0));

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "ER_Localize" command, which performs an Edmiston-Ruedenberg
///          localization on a given set of orbitals.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::ER_Localize_command_function(formic::Archive & arch, formic::AllData & data) {

  // read list of orbitals to localize
  std::vector<int> orb_inds;
  formic::mpi::read_and_bcast(arch, orb_inds, "Failed to read indices of orbitals to localize in formic::ER_Localize_command_function.");

  // read initial guess for local orbital coefficients
  std::vector<double> coeff_guess;
  formic::mpi::read_and_bcast(arch, coeff_guess, "Failed to read orbital coefficient guess in formic::ER_Localize_command_function.");

  // execute the localization
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::er_localize(data.userinp, data.hd, orb_inds, coeff_guess);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    throw formic::Exception("ER_Localize not implemented for complex numbers");
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"ER_Localize\"") % data.userinp.get<std::string>("scalar_type");

}
